package task_process

import (
	"context"
	"fmt"
	"shape-erver/utils"

	"github.com/gogf/gf/frame/g"
	"github.com/hibiken/asynq"
)

func AddShapeEToGIF(task_id, prompt string) error {
	fmt.Println("进到AddShapeEToGIF")
	return AddTaskToQueue(SHAPEE_SERVER_TASK, g.Map{"task_id": task_id, "prompt": prompt}, 0)
}

func ProcessShapeToGIF(c context.Context, t *asynq.Task) error {
	gpuId := -1
	avaliableGPUMap.Iterator(func(k int, v int) bool {

		if v == 1 { // 1 表示可用
			gpuId = k
			return false
		}
		return true
	})
	if gpuId == -1 {
		fmt.Println("没有可用的GPU")
		return nil
	} else {
		avaliableGPUMap.Set(gpuId, 0) // 0 表示不可用
	}

	fmt.Println("进到ProcessShapeToGIF")
	taskId, _ := t.Payload.GetString("task_id")
	prompt, _ := t.Payload.GetString("prompt")
	err := utils.DoRender(taskId, prompt, gpuId)
	if err != nil {
		avaliableGPUMap.Set(gpuId, 1) // 1 表示可用
		fmt.Println("AddToTaskQueue err:", err)
		return nil
	}
	avaliableGPUMap.Set(gpuId, 1) // 1 表示可用
	return nil
}

func CheckAvailableGPU() bool {
	gpuId := -1
	avaliableGPUMap.Iterator(func(k int, v int) bool {

		if v == 1 { // 1 表示可用
			gpuId = k
			avaliableGPUMap.Set(k, 0) // 0 表示不可用
			return false
		}
		return true
	})
	if gpuId == -1 {
		return false
	}

	return true
}
